import torch
import torch.nn as nn
from torch.optim import Adam
from dataset import _dataset
from torch.utils.data import DataLoader
import model

if __name__ == '__main__':
    test_dataset = _dataset("./datasets/test/")
    test_dataloader = DataLoader(test_dataset, batch_size=40, shuffle=True)

    train_dataset = _dataset("./datasets/train/")
    train_dataloader = DataLoader(train_dataset, batch_size=40, shuffle=True)

    _model = model.model1().cuda()
    # _model = torch.load("model.pth").cuda()
    loss_f = nn.MultiLabelSoftMarginLoss().cuda()
    optim = Adam(_model.parameters(), lr=0.001)

    for epoch in range(20):
        for i, (img, lab) in enumerate(train_dataloader):
            img = img.cuda()
            lab = lab.cuda()
            output = _model(img)
            loss = loss_f(output, lab)
            optim.zero_grad()
            loss.backward()
            optim.step()

        print("训练次数{} loss {}".format(epoch, loss.item()))

    torch.save(_model, "model.pth")